import pymysql
import pandas as pd
import datetime

class DatabaseConnector:
    def __init__(self):
        # 数据库连接配置
        self.connection = pymysql.connect(
            host="127.0.0.1",
            user="root",
            password="root",
            database="tushare",
            port=3306,
            charset="utf8",
            cursorclass=pymysql.cursors.DictCursor
        )

class FinancialAnalyzer:
    def __init__(self):
        pass

    def fetch_financial_indicators(self, connection):
        """获取财务指标数据并清洗"""
        with connection.cursor() as cursor:
            # 查询财务指标数据
            sql = """
            SELECT 
                ts_code, ann_date, eps, total_revenue_ps, undist_profit_ps, 
                gross_margin, fcff, fcfe, tangible_asset, bps, 
                grossprofit_margin, npta 
            FROM financial_d
            """
            cursor.execute(sql)
            results = cursor.fetchall()
        
        # 转换为DataFrame并清洗数据
        df = pd.DataFrame(results)
        cleaned_df = df.dropna(subset=[
            'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 
            'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta'
        ])
        cleaned_df = cleaned_df.reset_index(drop=True)
        
        return cleaned_df
    
    def analyze_stock_performance(self, connection, financial_data):
        """分析公告日后的股票表现"""
        records = []
        
        for idx, row in financial_data.iterrows():
            # 格式化公告日期
            announcement_date = row['ann_date'].strftime('%Y%m%d')
            
            # 查询公告日后20天的交易数据
            with connection.cursor() as cursor:
                sql = f"""
                SELECT trade_date, closes 
                FROM date_1 
                WHERE ts_code = '{row['ts_code']}' 
                AND trade_date > DATE('{announcement_date}') 
                ORDER BY trade_date ASC 
                LIMIT 20
                """
                cursor.execute(sql)
                trade_data = cursor.fetchall()
            
            # 处理交易数据
            try:
                if trade_data:
                    trade_df = pd.DataFrame(trade_data)
                    # 计算统计指标
                    max_price = trade_df['closes'].max()
                    min_price = trade_df['closes'].min()
                    second_day_price = trade_df['closes'].iloc[1]
                    
                    # 构建记录（修正了列引用错误）
                    records.append({
                        'ts_code': row['ts_code'],
                        'ann_date': row['ann_date'],
                        'max_close': max_price,
                        'min_close': min_price,
                        'the_close': second_day_price,
                        'eps': row['eps'],
                        'total_revenue_ps': row['total_revenue_ps'],  # 修正了行索引错误
                        'undist_profit_ps': row['undist_profit_ps'],  # 修正了行索引错误
                        'gross_margin': row['gross_margin'],
                        'fcff': row['fcff'],
                        'fcfe': row['fcfe'],
                        'tangible_asset': row['tangible_asset'],  # 修正了行索引错误
                        'bps': row['bps'],
                        'grossprofit_margin': row['grossprofit_margin'],
                        'npta': row['npta']
                    })
            except Exception as e:
                print(f"处理记录 {idx} 时出错: {e}")
            
            # 打印进度
            if idx % 100 == 0:
                print(f"已处理 {idx}/{len(financial_data)} 条记录")
        
        # 保存结果到CSV
        result_df = pd.DataFrame(records)
        result_df.to_csv('daily.csv', index=False)
        print(f"分析完成，结果已保存到 daily.csv，共 {len(result_df)} 条记录")

if __name__ == "__main__":
    # 创建数据库连接
    db_connector = DatabaseConnector()
    
    # 执行财务分析
    analyzer = FinancialAnalyzer()
    financial_data = analyzer.fetch_financial_indicators(db_connector.connection)
    
    # 注意：原代码未调用analyze_stock_performance方法，这里补充完整流程
    analyzer.analyze_stock_performance(db_connector.connection, financial_data)